from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.get_json import get_json

prompt = ChatPromptTemplate.from_template("""
你是评论员,你需要利用你的专业知识为提交文档提出一些有价值和意义的修改意见.
 - 需要评论的文档内容为:{submit}
 - 期撰写的主题为{title}
 - 未完成的段落为{sections}
 - 已通过的内容为{paragraphs}
 - 撰写者当前撰写段落主题为{section_current}
 - 你需要返回一个json对象,且只有两个key,第一个是pass(是否通过),第二个是idea(没有通过时,需要给出意见)
""")

judge_pass = {}


def generate_comment(state: EditorState):
    title = state['title']
    search = state['search']
    outlines = state['outlines']
    sections = outlines['sections']
    paragraphs = state['paragraphs']
    write_submit = state['write_submit']

    if f"{sections[0]}" not in judge_pass:
        judge_pass[f"{sections[0]}"] = 0
    judge_pass[f"{sections[0]}"] += 1

    section_current = sections[0]

    model = get_model()
    chains = prompt | model
    response = chains.invoke({"submit": write_submit, 'title': title, 'sections': sections, 'paragraphs': paragraphs,
                              "section_current": section_current})

    print("============ Comment Agent ============")
    print(response.content)
    comment_info = get_json(response.content)

    if judge_pass[f"{sections[0]}"] == 2:
        comment_info['pass'] = True
        outlines['sections'] = outlines['sections'][1:]

    who = "comment"
    if len(outlines['sections']) > 0:
        who = "outline"

    print(outlines['sections'])

    return {
        "who": who,
        "title": title,  # 制定标题
        "search": search,  # 搜集信息
        "outlines": outlines,  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 已完成的段落信息
        "write_submit": write_submit,
        "check_idea": {},  # 校对意见
        "comment_idea": comment_info  # 评论意见
    }


def comment_to_write_agent(state: EditorState):
    outlines = state['outlines']
    sections = outlines['sections']
    if len(sections) == 0 or sections == []:
        return "summary_agent"
    else:
        return "write_agent"

# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
